"""
文档读取策略模块

提供不同类型文档的读取策略和配置。
"""

from typing import Dict, List, Optional, Any
from abc import ABC, abstractmethod

from llama_index.core import SimpleDirectoryReader
from llama_index.core.schema import Document

from ..config import config


class DocumentReaderStrategy(ABC):
    """文档读取策略基类"""

    @abstractmethod
    def load_documents(self, file_paths: List[str], **kwargs) -> List[Document]:
        """
        加载文档

        Args:
            file_paths: 文件路径列表
            **kwargs: 额外参数

        Returns:
            List[Document]: 文档列表
        """
        pass


class LlamaIndexReaderStrategy(DocumentReaderStrategy):
    """基于LlamaIndex的读取策略"""

    def __init__(self):
        """初始化LlamaIndex读取策略"""
        self._readers = {}
        self._setup_readers()

    def _setup_readers(self) -> None:
        """设置不同文件类型的读取器"""
        # 注意：我们不在初始化时创建reader实例，而是在load_documents时动态创建
        # 这避免了SimpleDirectoryReader需要input_dir或input_files的问题
        self._reader_classes = {}

        # 标记哪些类型的reader可以使用
        self._available_readers = set(['txt'])  # txt总是可用

        # 尝试导入特定reader
        try:
            from llama_index.readers.file import PDFReader
            self._available_readers.add('pdf')
            self._reader_classes['pdf'] = PDFReader
        except ImportError:
            pass

        try:
            from llama_index.readers.file import DocxReader
            self._available_readers.add('docx')
            self._available_readers.add('doc')
            self._reader_classes['docx'] = DocxReader
            self._reader_classes['doc'] = DocxReader
        except ImportError:
            pass

        try:
            from llama_index.readers.file import MarkdownReader
            self._available_readers.add('md')
            self._available_readers.add('markdown')
            self._reader_classes['md'] = MarkdownReader
        except ImportError:
            pass

        try:
            from llama_index.readers.file import HTMLReader
            self._available_readers.add('html')
            self._available_readers.add('htm')
            self._reader_classes['html'] = HTMLReader
            self._reader_classes['htm'] = HTMLReader
        except ImportError:
            pass

        try:
            from llama_index.readers.file import PandasCSVReader
            self._available_readers.add('csv')
            self._reader_classes['csv'] = PandasCSVReader
        except ImportError:
            pass

  
    def load_documents(self, file_paths: List[str], **kwargs) -> List[Document]:
        """
        使用LlamaIndex加载文档

        Args:
            file_paths: 文件路径列表
            **kwargs: 额外参数

        Returns:
            List[Document]: 文档列表
        """
        if not file_paths:
            return []

        # 按文件扩展名分组
        grouped_files = self._group_files_by_type(file_paths)
        all_documents = []

        for file_type, paths in grouped_files.items():
            try:
                # 动态创建reader
                if file_type in self._available_readers:
                    if file_type == 'txt':
                        # 使用SimpleDirectoryReader处理文本文件
                        from llama_index.core.readers import SimpleDirectoryReader
                        reader = SimpleDirectoryReader(input_files=paths)
                        documents = reader.load_data()
                    else:
                        # 使用特定类型的reader
                        reader_class = self._reader_classes.get(file_type)
                        if reader_class:
                            documents = []
                            for file_path in paths:
                                try:
                                    doc = reader_class().load_data(file_path)
                                    if doc:
                                        documents.extend(doc)
                                except Exception as inner_e:
                                    print(f"读取文件 {file_path} 失败: {inner_e}")
                                    # 尝试使用SimpleDirectoryReader作为备用方案
                                    try:
                                        from llama_index.core.readers import SimpleDirectoryReader
                                        reader = SimpleDirectoryReader(input_files=[file_path])
                                        doc = reader.load_data()
                                        if doc:
                                            documents.extend(doc)
                                    except:
                                        continue
                        else:
                            # 使用SimpleDirectoryReader作为备用方案
                            from llama_index.core.readers import SimpleDirectoryReader
                            reader = SimpleDirectoryReader(input_files=paths)
                            documents = reader.load_data()
                else:
                    # 不支持的文件类型，使用SimpleDirectoryReader
                    from llama_index.core.readers import SimpleDirectoryReader
                    reader = SimpleDirectoryReader(input_files=paths)
                    documents = reader.load_data()

                all_documents.extend(documents)

            except Exception as e:
                print(f"读取 {file_type} 文件时出错: {e}")
                # 尝试使用SimpleDirectoryReader作为最后的备用方案
                try:
                    from llama_index.core.readers import SimpleDirectoryReader
                    reader = SimpleDirectoryReader(input_files=paths)
                    documents = reader.load_data()
                    all_documents.extend(documents)
                except Exception as backup_e:
                    print(f"备用方案也失败了: {backup_e}")
                continue

        return all_documents

    def _group_files_by_type(self, file_paths: List[str]) -> Dict[str, List[str]]:
        """按文件类型分组文件路径"""
        grouped = {}
        for file_path in file_paths:
            _, ext = file_path.lower().rsplit('.', 1)
            if ext not in grouped:
                grouped[ext] = []
            grouped[ext].append(file_path)
        return grouped


class EnhancedDocumentReader:
    """
    增强的文档读取器

    职责：
    1. 整合多种读取策略
    2. 提供统一的文档加载接口
    3. 处理读取错误和异常
    4. 支持批量处理和进度监控
    """

    def __init__(self, strategy: Optional[DocumentReaderStrategy] = None):
        """
        初始化增强文档读取器

        Args:
            strategy: 文档读取策略，默认使用LlamaIndex策略
        """
        self._strategy = strategy or LlamaIndexReaderStrategy()
        self._load_stats = {
            'total_files': 0,
            'successful_loads': 0,
            'failed_loads': 0,
            'failed_files': []
        }

    def load_documents(
        self,
        file_paths: List[str],
        show_progress: bool = False,
        max_retries: int = 3
    ) -> List[Document]:
        """
        加载文档

        Args:
            file_paths: 文件路径列表
            show_progress: 是否显示进度
            max_retries: 最大重试次数

        Returns:
            List[Document]: 加载的文档列表
        """
        if not file_paths:
            return []

        self._load_stats['total_files'] = len(file_paths)
        self._load_stats['successful_loads'] = 0
        self._load_stats['failed_loads'] = 0
        self._load_stats['failed_files'] = []

        documents = []

        for i, file_path in enumerate(file_paths):
            if show_progress and i % 10 == 0:
                print(f"正在处理文件 {i+1}/{len(file_paths)}: {file_path}")

            # 尝试加载单个文件
            file_docs = self._load_single_document_with_retry(
                file_path, max_retries
            )

            if file_docs:
                documents.extend(file_docs)
                self._load_stats['successful_loads'] += 1
            else:
                self._load_stats['failed_loads'] += 1
                self._load_stats['failed_files'].append(file_path)

        if show_progress:
            print(f"文档加载完成: 成功 {self._load_stats['successful_loads']}, "
                  f"失败 {self._load_stats['failed_loads']}")

        return documents

    def _load_single_document_with_retry(
        self, file_path: str, max_retries: int
    ) -> Optional[List[Document]]:
        """带重试机制的文档加载"""
        for attempt in range(max_retries):
            try:
                return self._strategy.load_documents([file_path])
            except Exception as e:
                if attempt == max_retries - 1:
                    print(f"加载文档失败 {file_path}: {e}")
                    return None
                # 简单的延迟重试
                import time
                time.sleep(0.5 * (attempt + 1))

        return None

    def get_load_statistics(self) -> Dict[str, Any]:
        """
        获取加载统计信息

        Returns:
            Dict[str, Any]: 统计信息
        """
        total = self._load_stats['total_files']
        successful = self._load_stats['successful_loads']
        failed = self._load_stats['failed_loads']

        return {
            'total_files': total,
            'successful_loads': successful,
            'failed_loads': failed,
            'success_rate': (successful / total * 100) if total > 0 else 0,
            'failed_files': self._load_stats['failed_files']
        }

    def reset_statistics(self) -> None:
        """重置统计信息"""
        self._load_stats = {
            'total_files': 0,
            'successful_loads': 0,
            'failed_loads': 0,
            'failed_files': []
        }

    def set_strategy(self, strategy: DocumentReaderStrategy) -> None:
        """设置读取策略"""
        self._strategy = strategy


# 默认文档读取器实例
default_reader = EnhancedDocumentReader()